﻿// vhd_bootSectUpdater.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"


#define MBR_MAGIC 0xAA55

struct CHS
{
	BYTE value1;
	BYTE value2;
	BYTE value3;

	WORD GetHead() const
	{
		return value1;
	}

	WORD GetCylinder() const
	{
		return value3 | ((value2 & 0x00c0) << 2);
	}

	WORD GetSector() const
	{
		return value2 & 0x003F;
	}
};

typedef struct 
{
	// 0x00	1	status (0x80 = bootable (active), 0x00 = non-bootable, other = invalid ) [a]
	BYTE	_reserved : 7;
	BYTE	active : 1;
	//0x01	3	CHS address of first absolute sector in partition. [b] The format is described by 3 bytes, see the next 3 rows.
	//0x01	1	h7-0	head [c]
	//0x02	1	c9‑8	s5–0	sector in bits 5–0; bits 7–6 are high bits of cylinder [c]
	//0x03	1	c7-0	bits 7–0 of cylinder [c]
	CHS		CHS_start;
	// 0x04	1	Partition type[7], http://en.wikipedia.org/wiki/Partition_type
	BYTE	type;
	//0x05	3	CHS address of last absolute sector in partition. [b] The format is described by 3 bytes, see the next 3 rows.
	//0x05	1	h7-0	head [c]
	//0x06	1	c9‑8	s5–0	sector in bits 5–0; bits 7–6 are high bits of cylinder [c]
	//0x07	1	c7-0	bits 7–0 of cylinder
	CHS		CHS_end;
	// 0x08	4	LBA of first absolute sector in the partition [d]
	DWORD   LBA_start;
	// 0x0C	4	Number of sectors in partition [d]
	DWORD   LBA_numOfSect;
}MBR_record,*PMBR_record;

enum FAT32_ATTRIB : BYTE
{
	FAT32_attr_readonly		= 1 << 0,	// (LSB)	Read Only	1	Should not allow writing
	FAT32_attr_hidden		= 1 << 1,	// Hidden				1	Should not show in dir listing
	FAT32_attr_system		= 1 << 2,	// System				1	File is operating system
	FAT32_attr_volumeId		= 1 << 3,	// Volume ID			1	Filename is Volume ID
	FAT32_attr_directory	= 1 << 4,	// Directory			x	Is a subdirectory (32-byte records)
	FAT32_attr_archive		= 1 << 5,	// Archive				x	Has been changed since last backup
	FAT32_attr_reserved1	= 1 << 6,	// Ununsed				0	Should be zero
	FAT32_attr_reserved2	= 1 << 7,	// (MSB)	Ununsed		0	Should be zero
};

typedef struct 
{
	BYTE name[11];
	FAT32_ATTRIB attrib;
	BYTE reserved;
	BYTE creationTimeLo;
	WORD creationTimeHi;	
	WORD creationDate;	
	WORD accessDate;	
	WORD firstClusterHi;
	WORD lastModTime;	
	WORD lastModDate;		
	WORD firstClusterLo;
	DWORD size;	
} FAT_DIR_ENTRY, *PFAT_DIR_ENTRY;

class FAT_DIR_ENTRY_CLUST
{
	PFAT_DIR_ENTRY _entries;

	WORD _entriesCount;

	WORD _currentEntry;
public:
	FAT_DIR_ENTRY_CLUST(LPVOID pFat32, DWORD firstClust = -1);

	PFAT_DIR_ENTRY GetNextEntry();

	void Reset();
};

typedef struct 
{
	BYTE bytes[64 * 1024];
	WORD size;
} BS_BUFFER, *PBS_BUFFER;

typedef struct 
{
	//0	 3	 The first three bytes 6B 3C and 90 disassemble to JMP SHORT 3C NOP. (The 3C value may be different.) The reason for this is to jump over the disk format information (the BPB and EBPB). Since the first sector of the disk is loaded into ram at location 0x0000:0x7c00 and executed, without this jump, the processor would attempt to execute data that isn't code.
	BYTE entry[3];
	//3	 8	 OEM identifier (mkdosfs). The first 8 Bytes (3 - 10) is the version of DOS being used. The next eight Bytes 29 3A 63 7E 2D 49 48 and 43 read out the name of the version. The official FAT Specification from Microsoft says that this field is really meaningless and is ignored by MS FAT Drivers, however it does recommend the value "MSWIN4.1" as some 3rd party drivers supposedly check it and expect it to have that value. Older versions of dos also report MSDOS5.1 and linux-formatted floppy will likely to carry "mkdosfs" here. If the string is less than 8 bytes, it is padded with zeros.
	BYTE OEM[8];
	//11	 2	 The number of Bytes per sector (remember, all numbers are in the little-endian format).
	// 512?
	WORD bytesPerSect;
	//13	 1	 Number of sectors per cluster.
	BYTE sectPerCluster;
	//14	 2	 Number of reserved sectors. The boot record sectors are included in this value.
	WORD reservedSectCount;
	//16	 1	 Number of File Allocation Tables (FAT's) on the storage media. Often this value is 2.
	BYTE fatCopiesAmount;
	//17	 2	 Number of directory entries (must be set so that the root directory occupies entire sectors).
	WORD dirEntries;
	//19	 2	 The total sectors in the logical volume. If this value is 0, it means there are more than 65535 sectors in the volume, and the actual count is stored in "Large Sectors (bytes 32-35).
	WORD sectPerVolume;
	//21	 1	 This Byte indicates the media descriptor type.
	BYTE mediaDescriptorType;
	//22	 2	 Number of sectors per FAT. FAT12/FAT16 only.
	WORD sectPerFAT1x;
	//24	 2	 Number of sectors per track.
	WORD sectPerTrack;
	//26	 2	 Number of heads or sides on the storage media.
	WORD heads;
	//28	 4	 Number of hidden sectors. (i.e. the LBA of the beginning of the partition.)
	DWORD hiddenSects;
	//32	 4	 Large amount of sector on media. This field is set if there are more than 65535 sectors in the volume.
	DWORD largeAmountOfSectors;
	//36	 4	 Sectors per FAT. The size of the FAT in sectors.
	DWORD sectPerFAT;
	//40	 2	 Flags.
	WORD flags;
	//42	 2	 FAT version number. The high byte is the major version and the low byte is the minor version. FAT drivers should respect this field.
	WORD version;
	//44	 4	 The cluster number of the root directory. Often this field is set to 2.
	DWORD clustOfRootDir;
	//48	 2	 The cluster number of the FSInfo structure.
	WORD clustOfFSInfo;
	//50	 2	 The cluster number of the backup boot sector.
	WORD clustOfBackupSector;
	//52	 12	 Reserved. When the volume is formated these bytes should be zero.
	BYTE reserved[12];
	//64	 1	 Drive number. The values here are identical to the values returned by the BIOS interrupt 0x13. 0x00 for a floppy disk and 0x80 for hard disks.
	BYTE driveNo;
	//65	 1	 Flags in windows NT. Reserved otherwise.
	BYTE ntFlags;
	//66	 1	 Signature (must be 0x28 or 0x29).
	BYTE signature;
	//67	 4	 VolumeID 'Serial' number. Used for tracking volumes between computers. You can ignore this if you want.
	DWORD serial;
	//71	 11	 Volume label string. This field is padded with spaces.
	BYTE label[11];
	//82	 8	 System identifier string. Always "FAT32   ". The spec says never to trust the contents of this string for any use.
	BYTE fsid[8];
	//90	 420	 Boot code.
	BYTE bootCode[420];
	//510	 2	 0xAA55 bootable partition signature.
	WORD magic;

	bool IsValid()
	{
		return magic == MBR_MAGIC &&
				(signature == 0x28 || signature == 0x29) &&
				memcmp(fsid, "FAT32   ", 8) == 0; 
	}

	bool Patch(LPVOID pHdd, PMBR_record pMbrRecord)
	{
		// Partition type must be one of 0x0B/0x0C/0x17/0x1B/0x1C/0x8B/0x8C/0x97/0x98/0xCB/0xCC
		switch (pMbrRecord->type)
		{
			case 0x0B: 
			case 0x0C: 
			case 0x17: 
			case 0x1B: 
			case 0x1C: 
			case 0x8B: 
			case 0x8C: 
			case 0x97: 
			case 0x98: 
			case 0xCB: 
			case 0xCC:
				break;
			default:
				return false;
		}

		PFAT32_BOOTSECTOR pFat32 = PFAT32_BOOTSECTOR(pMbrRecord->LBA_start * 512 + (PBYTE)pHdd);
		
		if (!pFat32->IsValid())
		{
			return false;
		}

		PBS_BUFFER pBS = (PBS_BUFFER)this;

		PBS_BUFFER pRealBS = (PBS_BUFFER)pFat32;

		memcpy(pFat32->bootCode, bootCode, sizeof(bootCode));
		memcpy(pFat32->entry, entry, sizeof(entry));

		if (pBS->size > 512)
		{
			memcpy(&pRealBS->bytes[512],
				   &pBS->bytes[512],
				   pBS->size - 512);
		}

		FSTest(pFat32);

		return true;
	}

	static void FSTest(LPVOID pFat32Buff)
	{
		PFAT32_BOOTSECTOR pFat32 = (PFAT32_BOOTSECTOR)pFat32Buff;

		if (pFat32->IsValid())
		{
			FAT_DIR_ENTRY_CLUST* dirReader = new FAT_DIR_ENTRY_CLUST(pFat32);
			
			PFAT_DIR_ENTRY pEntry = NULL;
			char name[12];

			name[11] = 0;

			for (;;)
			{
				pEntry = dirReader->GetNextEntry();

				if (!pEntry)
				{
					break;
				}

				memcpy(name, pEntry->name, 11);

				printf("%s\n\r", name);
				if (memcmp("XPRELDR SYS", pEntry->name, 11) == 0)
				{
					DWORD* clustChain = GetFileClustChain(pFat32, pEntry);
					//
					free(clustChain);
				}

			}
			
		}
	}

	static DWORD divide(DWORD numerator, DWORD denominator, DWORD* remainder)
	  {
		 *remainder = denominator % numerator;

		 return numerator  / denominator;
	  }

	static DWORD* GetFileClustChain(LPVOID pFat32Buff, PFAT_DIR_ENTRY pEntry)
	{
		PFAT32_BOOTSECTOR pFat32 = (PFAT32_BOOTSECTOR)pFat32Buff;

		
		DWORD clustAdd = 0;

		DWORD clustCount = divide(pEntry->size, (pFat32->sectPerCluster * pFat32->bytesPerSect), &clustAdd);
		
		if (clustAdd)
		{
			clustCount++;
		}

		DWORD* pFat = (DWORD*)((DWORD)pFat32 + pFat32->reservedSectCount * pFat32->bytesPerSect);

		DWORD clustNo = pEntry->firstClusterLo | (pEntry->firstClusterHi << 16);

		DWORD* clustChain = (DWORD*)calloc(clustCount, sizeof(DWORD));
		for (int i = 0; i < clustCount; i++)
		{
			clustChain[i] = clustNo;

			clustNo = pFat[clustNo];
		}

		return clustChain;
	}
} FAT32_BOOTSECTOR, *PFAT32_BOOTSECTOR;

FAT_DIR_ENTRY_CLUST::FAT_DIR_ENTRY_CLUST(LPVOID pBuff, DWORD firstClust /* = -1*/)
{
	PFAT32_BOOTSECTOR pFat32 = (PFAT32_BOOTSECTOR)pBuff;

	DWORD cluster_number = firstClust == -1 || firstClust < 2 ?
		pFat32->clustOfRootDir :
		firstClust;

	/*DWORD lengthOfTheFAT_inSectors = pFat32-> fatCopiesAmount * pFat32->sectPerFAT;
	DWORD positionOfTheBeginningOfData_inSectors= lengthOfTheFAT_inSectors + pFat32->reservedSectCount;

	_entries = (PFAT_DIR_ENTRY)(positionOfTheBeginningOfData_inSectors * pFat32->bytesPerSect
		+ (BYTE*)pFat32);*/

	_entriesCount = (pFat32->sectPerCluster * pFat32->bytesPerSect) / sizeof(FAT_DIR_ENTRY);
	_currentEntry = -1;

	unsigned long fat_begin_lba = (DWORD)pBuff + pFat32->reservedSectCount * pFat32->bytesPerSect;
	unsigned long cluster_begin_lba = fat_begin_lba + (pFat32->fatCopiesAmount * pFat32->sectPerFAT) * pFat32->bytesPerSect;

	unsigned char sectors_per_cluster = pFat32->sectPerCluster;
	unsigned long root_dir_first_cluster = pFat32->clustOfRootDir;

	
	_entries = (PFAT_DIR_ENTRY)(cluster_begin_lba + (cluster_number - 2) * sectors_per_cluster * pFat32->bytesPerSect);
}

PFAT_DIR_ENTRY FAT_DIR_ENTRY_CLUST::GetNextEntry()
{
getNextEntry:
	if (_currentEntry + 1 == _entriesCount)
	{
		return NULL;
	}

	PFAT_DIR_ENTRY entry = &_entries[(WORD)(_currentEntry + (WORD)1)];

	_currentEntry++;

	if (!entry->name[0])
	{
		goto getNextEntry;
	}

	if (entry->name[0] > 'Z' || entry->name[0] < ' ' ||
		(entry->attrib & FAT32_attr_volumeId) == FAT32_attr_volumeId ||
		(entry->attrib & (FAT32_attr_archive | FAT32_attr_hidden | FAT32_attr_system | FAT32_attr_readonly)) == (FAT32_attr_archive | FAT32_attr_hidden | FAT32_attr_system | FAT32_attr_readonly))
	{
		goto getNextEntry;
	}

	return entry;
}

void FAT_DIR_ENTRY_CLUST::Reset()
{
	_currentEntry = -1;
}

struct NTFS_BOOTSECTOR
{
//0x0B	Word	0x0002 (512)	Bytes per Sector
//0x0D	Byte	0x08	Sectors per Cluster
//0x0E	Word	0x000	Reserved Sectors
//0x10	3 Bytes	0x000000	reserved
//0x13	Word	0x0000	not used
//0x15	Byte	0xF8	Media Descriptor
//0x16	Word	0x0000	reserved
//0x18	Word	0x3F00 (63)	Sectors per Track
//0x1A	Word	0xFF00 (255)	Number of Heads
//0x1C	DWord	0x3F000000 (63)	Hidden Sectors
//0x20	DWord	0x00000000	not used
//0x24	DWord	0x00800080	not used
//0x28	QWord	0x00000000009C25FD	Total Sectors
//0x30	QWord	0x0000000000000004	$MFT Start Cluster
//0x38	QWord	0x000000000009C25F	$MFTmirr Start Cluster
//0x40	DWord	0x000000F6	Clusters per FRS
//0x44	DWord	0x00000001	Cluster per Index
//0x48	QWord	0xD83AEE5E98D83B12	Volume Serial Number
//0x50	DWord	0x00000000	CRC
	bool IsValid() const
	{
		return false;
	}
	
	bool Patch(LPVOID pHdd, PMBR_record pMbrRecord)
	{
		// Partition type must be one of 0x07/0x17/0x87
		return false;
	}
};

struct HFSPLUS_BOOTSECTOR
{
	bool IsValid() const
	{
		return false;
	}
	
	bool Patch(LPVOID pHdd, PMBR_record pMbrRecord)
	{
		// Partition type must be equals with 0xAF
		return false;
	}
};

enum FSTYPE
{
	FS_FAT32 = 0,
	FS_NTFS  = 0,
	FS_HFSP  = 0,
};

union BOOTSECTOR
{
	BS_BUFFER a;

	FAT32_BOOTSECTOR   fat32;
	NTFS_BOOTSECTOR    ntfs;
	HFSPLUS_BOOTSECTOR hfsplus;
};

typedef struct 
{
	BYTE code[0x01BE];
	/*
	0000	code area	440 (max. 446)
	01B8	disk signature (optional)	4
	01BC	Usually nulls; 0x0000	2
	01BE	Table of primary partitions
			(Four 16-byte entries, IBM partition table scheme)	64
	01FE	0776	510	55h	MBR signature	2
	01FF	0777	511	AAh
	*/

	MBR_record records[4];

	WORD	sign; //0xAA55
}MBR, *PMBR;


void UpdateBootRecord(LPVOID pHdd, BOOTSECTOR& bs)
{
	PMBR pMbr = (PMBR)pHdd;

	if (pMbr->sign != MBR_MAGIC)
	{
		MessageBox(NULL, _T("MBR code is invalid. I won't fix anything."), _T("Error"), MB_OK | MB_ICONEXCLAMATION);
		return;
	}

	PMBR_record pRecord = NULL;
	for (int i = 0; i < 4 ; i++)
	{
		if (pMbr->records[i].active)
		{
			pRecord = pMbr->records + i;
			break;
		}
	}

	if (pRecord)
	{
		if (pRecord->type)
		{
			bool patched = false;

			if (bs.fat32.IsValid())
			{
				patched = bs.fat32.Patch(pHdd, pRecord);
				if (!patched)
				{
					MessageBox(NULL, _T("Active partiotn is not valid FAT32 partition. No patch applied."), _T("Error"), MB_OK | MB_ICONASTERISK);
				}
			}
			else if (bs.ntfs.IsValid())
			{
				patched = bs.ntfs.Patch(pHdd, pRecord);
				if (!patched)
				{
					MessageBox(NULL, _T("Active partiotn is not valid NTFS partition. No patch applied."), _T("Error"), MB_OK | MB_ICONASTERISK);
				}
			}
			else if (bs.hfsplus.IsValid())
			{
				patched = bs.hfsplus.Patch(pHdd, pRecord);
				if (!patched)
				{
					MessageBox(NULL, _T("Active partiotn is not valid HFS+ partition. No patch applied."), _T("Error"), MB_OK | MB_ICONASTERISK);
				}
			}
	
			if (patched)
			{
				MessageBox(NULL, _T("Bootrecord was written successfully"), _T("Done"), MB_OK | MB_ICONINFORMATION);
			}
		}
		else
		{
			MessageBox(NULL, _T("Active partition has no FS type set"), _T("Error"), MB_OK | MB_ICONASTERISK);
		}
	}
	else
	{
		MessageBox(NULL, _T("No active partition was found"), _T("Error"), MB_OK | MB_ICONASTERISK);
	}
}

int _tmain(int argc, _TCHAR* argv[])
{
	BOOTSECTOR bs;
	ZeroMemory(bs.a.bytes, sizeof(BOOTSECTOR));

	bs.a.size = 0;

	OPENFILENAME ofn;
	TCHAR szFile[260];       // buffer for file name
	// Initialize OPENFILENAME
	ZeroMemory(&ofn, sizeof(ofn));
	ofn.lStructSize = sizeof(ofn);
	ofn.lpstrFile = szFile;
	// Set lpstrFile[0] to '\0' so that GetOpenFileName does not 
	// use the contents of szFile to initialize itself.
	ofn.lpstrFile[0] = '\0';
	ofn.nMaxFile = sizeof(szFile);
	ofn.lpstrFilter = _T("Binary bootsector (*.bin)\0*.bin");
	ofn.nFilterIndex = 1;
	ofn.lpstrFileTitle = NULL;
	ofn.nMaxFileTitle = 0;
	ofn.lpstrInitialDir = NULL;
	ofn.Flags = OFN_PATHMUSTEXIST | OFN_FILEMUSTEXIST;
	if (GetOpenFileName(&ofn))
	{
		HANDLE hBoot = CreateFile(ofn.lpstrFile, GENERIC_READ, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);

		if (hBoot == INVALID_HANDLE_VALUE)
		{
			MessageBox(NULL, _T("Selected file could not be opened"), _T("Error"), MB_OK | MB_ICONEXCLAMATION);
			return 0;
		}
		else
		{
			DWORD sz = GetFileSize(hBoot, NULL);
			DWORD bytesRead = 0;
			BOOL err = sz < 512 || sz > 64 * 1024 || !ReadFile(hBoot, &bs, sz, & bytesRead, NULL) || bytesRead < 512;
			if (err)
			{
				MessageBox(NULL, _T("Input file does not contain valid boot sector or couldn't be read."), _T("Error"), MB_OK | MB_ICONEXCLAMATION);
			}
			else
			{
				bs.a.size = sz;

				if (!bs.fat32.IsValid() && 
					!bs.hfsplus.IsValid() &&
					!bs.ntfs.IsValid())
				{
					MessageBox(NULL, _T("Bootsector contains information for unknown filesystem. I will do nothing!"), _T("Error"), MB_OK | MB_ICONEXCLAMATION);

					err = TRUE;
				}
			}

			CloseHandle(hBoot);

			if (err)
			{
				return 0;
			}
		}
	}
	else
	{
		MessageBox(NULL, _T("No bootsector was chosen"), _T("Error"), MB_OK | MB_ICONEXCLAMATION);
		return 0;
	}


	// Initialize OPENFILENAME
	ZeroMemory(&ofn, sizeof(ofn));
	ofn.lStructSize = sizeof(ofn);
	ofn.lpstrFile = szFile;
	// Set lpstrFile[0] to '\0' so that GetOpenFileName does not 
	// use the contents of szFile to initialize itself.
	ofn.lpstrFile[0] = '\0';
	ofn.nMaxFile = sizeof(szFile);
	ofn.lpstrFilter = _T("Virtual Hard Disk\0*.vhd");
	ofn.nFilterIndex = 1;
	ofn.lpstrFileTitle = NULL;
	ofn.nMaxFileTitle = 0;
	ofn.lpstrInitialDir = NULL;
	ofn.Flags = OFN_PATHMUSTEXIST | OFN_FILEMUSTEXIST;


	if (GetOpenFileName(&ofn))
	{
		HANDLE hFile = CreateFile( ofn.lpstrFile, GENERIC_READ | GENERIC_WRITE, FILE_SHARE_READ, NULL, OPEN_EXISTING, 0, NULL);

		if (hFile == INVALID_HANDLE_VALUE)
		{
			MessageBox(NULL, _T("File could not be opened"), _T("Error"), MB_OK | MB_ICONEXCLAMATION);
		}
		else
		{
			DWORD sz = GetFileSize(hFile, NULL);
			DWORD LBA = sz / 512;

			BYTE heads_per_cylinder = 0x10;
			BYTE sectors_per_track = 0x3f;

			DWORD cylinder = LBA / (heads_per_cylinder * sectors_per_track);
	        DWORD temp = LBA % (heads_per_cylinder * sectors_per_track);
	        DWORD head = temp / sectors_per_track;
	        DWORD sector = temp % sectors_per_track + 1;

			printf("Bytes: %d\n"
				   "LBA: 0x%08x\n"
				   "CHS: %d:%d:%d\n\n", sz, LBA, cylinder, head, sector);

			HANDLE hMap = CreateFileMapping(hFile, NULL, PAGE_READWRITE, 0, 0, NULL);

			if (!hMap)
			{
				MessageBox(NULL, _T("Mapping could not be created"), _T("Error"), MB_OK | MB_ICONEXCLAMATION);
			}
			else
			{
				LPVOID pHdd = MapViewOfFile(hMap, FILE_MAP_ALL_ACCESS, 0, 0, 0);

				if (pHdd)
				{
					UpdateBootRecord(pHdd, bs);

					UnmapViewOfFile(pHdd);
				}
				else
				{
					MessageBox(NULL, _T("OS was unable to map VHD to memoryspace"), _T("Error"), MB_OK | MB_ICONEXCLAMATION);
				}

				CloseHandle(hMap);
			}

			CloseHandle(hFile);
		}
	}

	return 0;
}

